`include "mycpu.h"

module mem_stage(
    input                          clk           ,
    input                          reset         ,
    //allowin
    input                          ws_allowin    ,
    output                         ms_allowin    ,
    //from es
    input                          es_to_ms_valid,
    input  [`ES_TO_MS_BUS_WD -1:0] es_to_ms_bus  ,
    //to ws
    output                         ms_to_ws_valid,
    output [`MS_TO_WS_BUS_WD -1:0] ms_to_ws_bus  ,
    //from data-sram
    input  [31                 :0] data_sram_rdata,
    //add block and bypass code
    output [37                 :0] ms_to_id_byp_block,
    //Int
    output [0                  :0] ms_ex,
    output [0                  :0] ms_ertn,
    output [0                  :0] ms_inst_csrrd,
    output [0                  :0] ms_inst_csrwr,
    output [0                  :0] ms_inst_csrxchg,
    input  [0                  :0] ws_ertn,
    input  [0                  :0] ws_ex  
);

assign ms_inst_csrrd = inst_csrrd & ms_valid;
assign ms_inst_csrwr = inst_csrwr & ms_valid;
assign ms_inst_csrxchg = inst_csrxchg & ms_valid;

reg         ms_valid;
wire        ms_ready_go;

reg [`ES_TO_MS_BUS_WD -1:0] es_to_ms_bus_r;
wire        ms_res_from_mem;
wire        ms_gr_we;
wire [ 4:0] ms_dest;
wire [31:0] ms_alu_result;
wire [31:0] ms_pc;
wire [2:0]  ms_store_op;
wire [4:0]  ms_load_op;

//Lab 8 add code, es_to_ms bus ADD signals!!!!
wire [31:0] ms_rj_value;
wire [31:0] ms_rkd_value;
wire [31:0] ms_inst;

wire [0: 0] inst_csrrd;
wire [0: 0] inst_csrwr;
wire [0: 0] inst_csrxchg;
wire [0: 0] inst_ertn;
wire [0: 0] inst_syscall;

wire [5: 0] ms_ecode;
wire [8: 0] ms_esubcode;
wire [0: 0] ms_has_int;

assign {ms_esubcode    ,  //196:188
        ms_ecode       ,  //187:182
        inst_syscall   ,  //181:181
        ms_has_int     ,  //180:180
        ms_ex          ,  //179:179
        inst_ertn      ,  //178:178
        ms_inst        ,  //177:146
        inst_csrxchg   ,  //145:145
        inst_csrwr     ,  //144:144
        inst_csrrd     ,  //143:143
        ms_rkd_value   ,  //142:111
        ms_rj_value    ,  //110:79
        ms_load_op     ,  //78:74
        ms_store_op    ,  //73:71
        ms_res_from_mem,  //70:70
        ms_gr_we       ,  //69:69
        ms_dest        ,  //68:64
        ms_alu_result  ,  //63:32
        ms_pc             //31:0
       } = es_to_ms_bus_r;

assign ms_ertn = inst_ertn & ms_valid;

wire [31:0] mem_result;
wire [31:0] ms_final_result;

//We need to add signals in the ms_to_ws_bus
assign ms_to_ws_bus = {ms_esubcode    ,  //187:179
                       ms_ecode       ,  //178:173
                       inst_syscall   ,  //172:172
                       ms_has_int     ,  //171:171
                       ms_ex          ,  //170:170
                       inst_ertn      ,  //169:169
                       ms_inst        ,  //168:137
                       inst_csrxchg   ,  //136:136
                       inst_csrwr     ,  //135:135
                       inst_csrrd     ,  //134:134
                       ms_rkd_value   ,  //133:102
                       ms_rj_value    ,  //101:70
                       ms_gr_we       ,  //69:69
                       ms_dest        ,  //68:64
                       ms_final_result,  //63:32
                       ms_pc             //31:0
                      };

assign ms_ready_go    = 1'b1;
assign ms_allowin     = !ms_valid || ms_ready_go && ws_allowin;
assign ms_to_ws_valid = ms_valid && ms_ready_go && ~ws_ertn && ~ws_ex;
always @(posedge clk) begin
    if (reset) begin
        ms_valid <= 1'b0;
    end
    else if (ms_allowin) begin
        ms_valid <= es_to_ms_valid;
    end

    if (es_to_ms_valid && ms_allowin) begin
        es_to_ms_bus_r  <= es_to_ms_bus;
    end
end


//Lab 7 change code
wire [31:0] lb_result;
wire [31:0] lbu_result;
wire [31:0] lh_result;
wire [31:0] lhu_result;

assign lb_result = (ms_alu_result[1:0] == 2'b00)? {{24{data_sram_rdata[7]}},data_sram_rdata[7:0]}:
                   (ms_alu_result[1:0] == 2'b01)? {{24{data_sram_rdata[15]}},data_sram_rdata[15:8]}:
                   (ms_alu_result[1:0] == 2'b10)? {{24{data_sram_rdata[23]}},data_sram_rdata[23:16]}:
                                                  {{24{data_sram_rdata[31]}},data_sram_rdata[31:24]}  ;
assign lbu_result = (ms_alu_result[1:0] == 2'b00)? {24'b0,data_sram_rdata[7:0]}:
                    (ms_alu_result[1:0] == 2'b01)? {24'b0,data_sram_rdata[15:8]}:
                    (ms_alu_result[1:0] == 2'b10)? {24'b0,data_sram_rdata[23:16]}:
                                                   {24'b0,data_sram_rdata[31:24]}  ;
assign lhu_result = (ms_alu_result[1] == 1'b0)? {16'b0,data_sram_rdata[15:0]}:{16'b0,data_sram_rdata[31:16]};   
assign lh_result = (ms_alu_result[1] == 1'b0)? {{16{data_sram_rdata[15]}},data_sram_rdata[15:0]}:
                                               {{16{data_sram_rdata[31]}},data_sram_rdata[31:16]};   


assign mem_result = ms_load_op[4] ? data_sram_rdata:
                    ms_load_op[3] ? lb_result      :
                    ms_load_op[2] ? lh_result:
                    ms_load_op[1] ? lbu_result: lhu_result;
                

assign ms_final_result = ms_res_from_mem ? mem_result
                                         : ms_alu_result;

//add block and bypass code
wire [4:0] ms_to_id_dest;
wire [0:0] ms_wen;
wire [31:0] ms_to_id_res;

assign ms_to_id_dest = ms_dest;
assign ms_wen        = ms_gr_we && ms_valid;
assign ms_to_id_res  = ms_alu_result;

assign ms_to_id_byp_block = {ms_wen, ms_to_id_dest,ms_final_result};

endmodule
